#!/usr/bin/env python
# coding: utf-8

from itertools import product, groupby

ALPHABET = ['a','c','t','g']

def compl(char):
    if char == 'a': return 'g'
    if char == 'g': return 'a'
    if char == 'c': return 't'
    if char == 't': return 'c'
    

class SimpleCreator(object):
    """
    generate()
        Для множества пар с весами (t1,t2) делаем расчёт, сохраняем результат для ключей keys
    Использование:
        Для входной пары (t1,t2) определяем ключ key, с ним идём в БД.
    """
    def __init__(self, n):
        self.points_num = n
    
    def _create_key(self, weights1, weights2):
        assert weights1.keys() == weights2.keys() == ALPHABET
        key = str(self.points_num) + '|' + ','.join('%i' % min(weights1[c], weights2[c]) for c in ALPHABET)
        return key
    
    def generate(self,):
        """
        keys - ключи в базе данных, для которых результат рассчёта для пары (t1,t2) одинаков
        t1, t2 - туплы весов символов двух последовательностей
        """
        l = range(self.points_num+1)
        ll=product(l, repeat=len(ALPHABET))
        llx = [i for i in ll if sum(i) == self.points_num]
        llxy = list(product(llx, repeat=2))
        
        llxyz = {}
        for t1,t2 in llxy:
            w1 = dict(zip(ALPHABET, t1))
            w2 = dict(zip(ALPHABET, t2))
            key = self._create_key(w1, w2)
            llxyz[key] = tuple(sorted((t1,t2)))
        
        for (t1,t2), grp in groupby(sorted(llxyz.iteritems(), key=lambda i:i[1]), key=lambda i:i[1]):
            keys = [key for key,t in grp]
            t1 = tuple(float(t)/self.points_num for t in t1)
            t2 = tuple(float(t)/self.points_num for t in t2)
            w1 = dict(zip(ALPHABET, t1))
            w2 = dict(zip(ALPHABET, t2))
            yield (w1,w2), keys
    
    def _adjust_weights(self, weights):
        diff = sum(weights.values()) - self.points_num
        while diff:
            if diff > 0:
                c = sorted(weights.items(), key=lambda i:i[1])[-1][0]
                weights[c] -= 1
            else:
                c = sorted(filter(lambda (c,w): w>0, weights.items()), key=lambda i:i[1])[0][0]
                weights[c] += 1
            diff = sum(weights.values()) - self.points_num
    
    def get_key(self, d1, d2):
        assert sum(d1.values()) == 1.0
        assert sum(d2.values()) == 1.0
        weights1 = dict((c, int(round(w*self.points_num))) for c,w in d1.items())
        weights2 = dict((c, int(round(w*self.points_num))) for c,w in d2.items())
        
        if sum(weights1.values()) != self.points_num:
            self._adjust_weights(weights1)
        if sum(weights2.values()) != self.points_num:
            self._adjust_weights(weights2)
        
        return self._create_key(weights1, weights2)



class PuPyCreator(SimpleCreator):
    def _create_key(self, weights1, weights2):
        assert weights1.keys() == weights2.keys() == ALPHABET
        key = str(self.points_num) + '|' + ','.join(
            '%i:%i' % (
                min(weights1[c], weights2[c]),
                min(weights1[c], weights2[compl(c)]),
            )
            for c in ALPHABET
        )
        return key


def format_generation(w1, w2, keys):
    return '%s\t%s\t%s' % (
        ','.join('%s:%g' % item for item in w1.items()),
        ','.join('%s:%g' % item for item in w2.items()),
        ' '.join(keys),
    )

def parse_generation(s):
    spl = s.split('\t')
    w1  = dict((c, float(w)) for c,w in (cw.split(':') for cw in spl[0].split(',')))
    w2  = dict((c, float(w)) for c,w in (cw.split(':') for cw in spl[1].split(',')))
    keys = spl[2].split(' ')
    return w1, w2, keys



DEGREES = 5

def interpolate_learn(_x, _y):
    import numpy as np
    x, y = np.array(_x), np.array(_y)
    coeffs = np.polyfit(x, y, DEGREES)
    return list(reversed(coeffs))

def interpolate_apply(coeffs):
    def apply(x):
        res = 0
        for i, coef in enumerate(coeffs):
            res += coef * x**i
        return res
    return apply


